package com.zyg.javaagent;

import com.sun.org.apache.bcel.internal.generic.IF_ACMPEQ;
import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

/**
 * 字节马插桩类
 * 入口静态方法 ：premain(String arg,Instrumentation instrumentation)
 * 打包参数入口类：premain-class
 * 运行参数：需要依赖其他程序运行 -javaagent:xxx.jar
 */
public class AgentMain {
    /**
     * javaagent入口方法
     * @param arg
     * @param instrumentation
     */
    //1.1 insertBefore和insertAfter,底层采用代码块的方式插入，可能存在变量无法获取问题
//    public static void premain(String arg, Instrumentation instrumentation) {
//        System.out.println("hello javaagent");
//        //基于javassist工具 在运行时修改class字节码，即插桩逻辑
//        final ClassPool pool = new ClassPool();
//        pool.appendSystemPath();
//        //添加类加载过滤器
//        instrumentation.addTransformer(new ClassFileTransformer() {
//            @Override
//            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
//                try {
//                    if (!"com/zyg/javaagent/UserService".equals(className)) {
//                        //不是该类不做修改
//                        return null;
//                    }
//                    CtClass ctClass = pool.get("com.zyg.javaagent.UserService");
//                    CtMethod sayHello = ctClass.getDeclaredMethod("sayHello");
//                    //底层采用代码块的方式插入
//                    sayHello.insertBefore("System.out.println(System.currentTimeMillis());");
//                    return ctClass.toBytecode();
//                } catch (Exception e) {
//                    e.printStackTrace();
//                }
//                return null;
//            }
//        });
//    }
//    //1.2 插桩sayHello方法(无参数)
//    public static void premain(String arg, Instrumentation instrumentation) {
//        System.out.println("hello javaagent");
//        //基于javassist工具 在运行时修改class字节码，即插桩逻辑
//        final ClassPool pool = new ClassPool();
//        pool.appendSystemPath();
//        //添加类加载过滤器
//        instrumentation.addTransformer(new ClassFileTransformer() {
//            @Override
//            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
//                try {
//                    if (!"com/zyg/javaagent/UserService".equals(className)) {
//                        //不是该类不做修改
//                        return null;
//                    }
//                    CtClass ctClass = pool.get("com.zyg.javaagent.UserService");
//                    CtMethod sayHello = ctClass.getDeclaredMethod("sayHello");
//                    //复制方法
//                    CtMethod copy = CtNewMethod.copy(sayHello, ctClass, null);
//                    copy.setName("sayHello$agent");
//                    ctClass.addMethod(copy);
//                    //改变原有方法
//                    sayHello.setBody("{ long begin = System.currentTimeMillis();\n" +
//                            "        sayHello$agent();\n" +
//                            "        System.out.println(System.currentTimeMillis()-begin); }");
//                    return ctClass.toBytecode();
//                } catch (Exception e) {
//                    e.printStackTrace();
//                }
//                return null;
//            }
//        });
//    }

//    //1.2 插桩print方法（带参数）
//    public static void premain(String arg, Instrumentation instrumentation) {
//        System.out.println("hello javaagent");
//        //基于javassist工具 在运行时修改class字节码，即插桩逻辑
//        final ClassPool pool = new ClassPool();
//        pool.appendSystemPath();
//        //添加类加载过滤器
//        instrumentation.addTransformer(new ClassFileTransformer() {
//            @Override
//            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
//                try {
//                    if (!"com/zyg/javaagent/UserService".equals(className)) {
//                        //不是该类不做修改
//                        return null;
//                    }
//                    CtClass ctClass = pool.get("com.zyg.javaagent.UserService");
//                    CtMethod sayHello = ctClass.getDeclaredMethod("print");
//                    //复制方法
//                    CtMethod copy = CtNewMethod.copy(sayHello, ctClass, null);
//                    copy.setName("print$agent");
//                    ctClass.addMethod(copy);
//                    //改变原有方法
//                    sayHello.setBody("{ long begin = System.currentTimeMillis();\n" +
//                            "        print$agent($1);\n" +
//                            "        System.out.println(System.currentTimeMillis()-begin); }");
//                    return ctClass.toBytecode();
//                } catch (Exception e) {
//                    e.printStackTrace();
//                }
//                return null;
//            }
//        });
//    }

    //1.3 通用方法插桩

    /**
     * 监控指定区域的类或方法    如指定包下的类com.zyg.javaagent.*
     * 1.读取参数
     * 2.匹配类
     * 3.获取类中的所有方法
     * 4.改造方法
     *
     * @param arg
     * @param instrumentation
     */
    public static void premain(String arg, Instrumentation instrumentation) {
        System.out.println("hello javaagent");
        String config = arg;
        //基于javassist工具 在运行时修改class字节码，即插桩逻辑
        final ClassPool pool = new ClassPool();
        pool.appendSystemPath();
        //添加类加载过滤器
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (className == null || !className.replaceAll("/", ".").startsWith(config)) {
                    return null;
                }
                className = className.replaceAll("/", ".");
                try {
                    CtClass ctClass = pool.get(className);
                    //获取类中的所有方法
                    for (CtMethod declaredMethod : ctClass.getDeclaredMethods()) {
                        newMethod(declaredMethod);
                    }
                    return ctClass.toBytecode();
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return null;
            }
        });
    }

    private static void newMethod(CtMethod oldMethod) throws Exception {
        CtClass ctClass = oldMethod.getDeclaringClass();
        CtMethod copy = CtNewMethod.copy(oldMethod, ctClass, null);
        copy.setName(oldMethod.getName() + "$agent");
        ctClass.addMethod(copy);
        //改变原有方法
        if (oldMethod.getReturnType().equals(CtClass.voidType)) {
            oldMethod.setBody(String.format(voidSource, oldMethod.getName(),oldMethod.getName()));
        } else {
            oldMethod.setBody(String.format(notVoidSource, oldMethod.getName(),oldMethod.getName()));
        }
    }

    //无返回值的方法改造模式
    static String voidSource = "{" +
            "long begin = System.currentTimeMillis();\n" +
            "        try {\n" +
            "            %s$agent($$);\n" +
            "        } finally {\n" +
            "            long end = System.currentTimeMillis();\n" +
            "            System.out.println(\"%s方法执行时间为\"+(end-begin));\n" +
            "        }}";
    //有返回值的方法改造模式
    static String notVoidSource = "{" +
            "        long begin = System.currentTimeMillis();\n" +
            "        Object result ;\n" +
            "        try {\n" +
            "            result = ($w) %s$agent($$);\n" +
            "        } finally {\n" +
            "            long end = System.currentTimeMillis();\n" +
            "            System.out.println(\"%s方法执行时间为\"+(end-begin));\n" +
            "        }\n" +
            "        return ($r) result;" +
            "}";
}
